#include "gemm.h"

template<typename T>
__global__ void cuda_gemm(bool transpose_l,bool transpose_r,int m,int n,int K,T* A,T* B,T* C,int rows,int cols,int total_num){
    int idx = blockIdx.x*blockDim.x + threadIdx.x;
    if (idx < total_num) {
        int row_idx = idx / cols;
        int col_idx = idx % cols;
        T temp = 0;
        for (int k = 0; k < K; ++k) {
            if (!transpose_l && !transpose_r) { // A * B
                temp += A[row_idx * K + k] * B[k * cols + col_idx];
            }
            else if (!transpose_l && transpose_r) { // A * B.T
                temp += A[row_idx * K + k] * B[col_idx * K + k];
            }
            else if (transpose_l && !transpose_r) { // A.T * B
                temp += A[k * rows + row_idx] * B[k * cols + col_idx];
            }
            else { // A.T * B.T
                temp += A[k * rows + row_idx] * B[col_idx * K + k];
            }
        }
        C[row_idx * cols + col_idx] = temp;
    }
}

// template<>
// __global__ void cuda_gemm(bool tranpose_l,bool transpose_r,half* A,half* B,half* C,int rows,int cols,int total_num){
    
// }


template<typename T>
void launch_cuda_gemm(bool tranpose_l,bool transpose_r,T* A,T* B,T* C,int rows,int cols,int blocksize,int total_nums){
    int gridesize = (n+blocksize-1)/blocksize;
    cuda_gemm<<<gridesize,blocksize>>>(tranpose_l,transpose_r,A,B,C,rows,cols,total_nums);  

}


// template<typename T>
// void launch_cpu_gemm(bool tranpose_l,bool transpose_r,T* A,T* B,T* C,int rows,int cols);

// template<>
// void launch_cpu_gemm(bool tranpose_l,bool transpose_r,half* A,half* B,half* C,int rows,int cols);



template void launch_cuda_gemm<float>(bool,bool,float*,float*,float*,int,int,int,int);
// template void launch_cuda_matmul<half>(half*,half*,half*,int,int,int,int);
// template void launch_cuda_matmul<uint8_t>(uint8_t*,uint8_t*,uint8_t*,int,int,int,int);

// template void launch_cpu_matmul<float>(float*,float*,float*,int,int);
// template void launch_cpu_matmul<half>(half*,half*,half*,int,int);
// template void launch_cpu_matmul<uint8_t>(uint8_t*,uint8_t*,uint8_t*,int,int);